/*
 * topological-hex, a program to compute combinatorial hexahedral meshes.
 *
 * Copyright (C) <2018> <Université catholique de Louvain (UCL), Belgique>
 *
 * List of the contributors to the development, description and complete
 * License: see LICENSE file.
 *
 * This program (topological-hex) is free software:
 * you can redistribute it and/or modify it under the terms
 * of the GNU General Public License as published by the Free
 * Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program (see COPYING file).  If not,
 * see <http://www.gnu.org/licenses/>.
 */

#include "server.h"
#include "sparse_set.h"

#include <limits.h>
#include <stdlib.h>
#include <unistd.h>
#include <string.h>
#include <stdarg.h>
#include <stdio.h>
#include <fcntl.h>

#include <sys/socket.h>
#include <sys/types.h>
#include <netdb.h>
#include <arpa/inet.h>
#include <errno.h>

#include <pthread.h>

/**
 * Stores all information used when processing clients.
 */
typedef struct client_data {
  /**
   * C-string used to identify the client in log files.
   */
  const char *name;

  /**
   * Unique integer index used to identify the client.
   */
  uint64_t id;

  int socket;

  /**
   * Index of the last subproblem visited by the algorithm.
   */
  uint32_t cursor_id;

  /**
   * Buffer large enough for any packet that will be sent to or received from
   * the client.
   */
  void *buffer;

  /**
   * Set of subproblems the client is currently exploring.
   */
  sparse_set exploring;

  /**
   * Capacity of the @ref client_data::id_buffer array.
   */
  uint32_t id_capacity;

  /**
   * Array used to store the indices of the subproblems that will be sent to the
   * client.
   */
  uint32_t *id_buffer;

  /**
   * Whether or not the client is currently allowed to send requests to the
   * server.
   */
  bool asleep;
} client_data;

/**
 * Data-structures shared by all threads to maintain the state of the server
 * during a distributed computation.
 */
typedef struct server {
  int socket;

  server_callback cb;
  void *user_data;

  server_log log;
  void *log_data;

  server_listen_cb listen_cb;
  void *listen_data;

  server_progress_cb progress_cb;
  void *progress_data;

  /**
   * If set to true, has received or is currently receiving a layout for its
   * subproblems.
   */
  atomic_bool has_layout;

  /**
   * If true, the server has received its initial set of subproblems.
   */
  bool has_started;

  /**
   * Set to true to indicate that clients should be disconnected.
   */
  bool should_close;

  /**
   * Layout used to represent subproblems.
   */
  subproblem_layout layout;

  /**
   * Memory buffer containing all subproblems sent by the clients, in their raw
   * form (i.e. the same format used to transmit them).
   *
   * The size of each element is computed is @ref
   * subproblem_layout::packet_size.
   */
  uint8_t *subproblems;

  /**
   * Index of the current thread exploring each subproblem.
   */
  atomic_uint_fast32_t *owners;

  /**
   * Number of subproblems stored in @ref server::subproblems.
   */
  uint32_t num_subproblems;

  /**
   * Maximum number of subproblems stored in @ref server::subproblems.
   */
  uint32_t capacity;

  /**
   * Number of clients that have connected to the server.
   */
  uint64_t num_clients;

  /**
   * Number of subproblems that have not yet been sent to any client.
   */
  atomic_uint_fast32_t num_available;

  /**
   * Number of subproblems that still need to be explored.
   */
  atomic_uint_fast32_t num_unfinished;

  /**
   * Threads used to handle communication with clients.
   */
  pthread_t *threads;

  /**
   * IDs of the clients handled by threads stored in @ref server::threads.
   */
  uint64_t *client_ids;

  /**
   * Number of clients which are being handled by a thread still alive.
   */
  uint32_t num_active_threads;

  /**
   * Capacity of the @ref server::threads array.
   */
  uint32_t thread_capacity;

  /**
   * Lock that should be acquired when manipulating @ref server::threads.
   */
  pthread_mutex_t threads_lock;

  /**
   * If set to true, more problems have been finished since the last time the
   * backup thread wrote the bitmap.
   */
  atomic_bool bitmap_dirty;

  /**
   * File descriptor to write solutions to. This is used to restore the state of
   * the program in the event of an interruption.
   *
   * This value is set to -1 if backups should not be maintained by the
   * algorithm.
   */
  int solutions_fd;

  /**
   * Lock that should be acquired when writing a solution to a file.
   */
  pthread_mutex_t solutions_lock;
} server;

/**
 * Information about a new client
 */
typedef struct client_parameters {
  server *server;

  int socket;
  struct sockaddr_storage addr;

  uint64_t id;
} client_parameters;

/**
 * Information used to store the progress of the program on a file system. This
 * allows minimal progress to be lost if the server needs to be restarted.
 */
typedef struct backup_parameters {
  server *server;

  /**
   * These buffers must have a capacity of at least 1024 bytes.
   */
  char *bitmap_path;
  char *temp_path;
  char *subproblems_path;
  char *solutions_path;

  bool try_restore;
} backup_parameters;

/**
 * Logs a message, which can use printf-style formatting commands.
 */
static void server_printf(server *server, const char *fmt, ...) {
  char buffer[1024];
  va_list list;
  va_start(list, fmt);
  vsnprintf(buffer, sizeof(buffer), fmt, list);
  va_end(list);

  server->log(buffer, server->log_data);
}

/**
 * Initial amount of subproblems that should be allocated when waiting for a
 * connection.
 */
#define SERVER_CAPACITY (8192)

#define OWNER_DONE ((uint32_t)-1)

/**
 * Initializes memory needed for the server to run.
 */
static error_code server_init(server *server);

/**
 * Releases all memory used by the server.
 */
static void server_release(server *server);

/**
 * Starts listening for connections.
 */
static error_code server_open(server *server);

/**
 * Tries to accept a new client.
 */
static error_code server_accept(server *server);

static error_code spawn_backup_thread(server *server, pthread_t *backup_thread,
                                      server_backup_info *backup,
                                      backup_parameters *params);

static error_code spawn_client_thread(
  server *server, client_parameters *params);

static void *client_handler(void *arg);

static void *backup_handler(void *arg);

error_code server_run(server_callback cb, void *user_data,
                      server_log log, void *log_data,
                      server_listen_cb listen_cb, void *listen_data,
                      server_progress_cb progress, void *progress_data,
                      server_backup_info *backup) {
  server server;

  server.cb = cb;
  server.user_data = user_data;

  server.log = log;
  server.log_data = log_data;

  server.listen_cb = listen_cb;
  server.listen_data = listen_data;

  server.progress_cb = progress;
  server.progress_data = progress_data;

  error_code status = SUCCESS;

  if ((status = server_init(&server)) != SUCCESS)
    return status;

  char bitmap_file[1024];
  char temp_file[1024];
  char subproblems_file[1024];
  char solutions_file[1024];

  backup_parameters backup_params = {
    &server,
    bitmap_file, temp_file, subproblems_file, solutions_file,
    backup ? backup->try_restore : false
  };

  pthread_t backup_thread;

  if (backup) {
    if ((status = spawn_backup_thread(&server, &backup_thread,
                                      backup, &backup_params)) != SUCCESS) {
      server_release(&server);
      return status;
    }
  }

  if ((status = server_open(&server)) != SUCCESS) {
    if (backup) pthread_join(backup_thread, NULL);
    server_release(&server);
    return status;
  }

  server.listen_cb(server.listen_data);
  server_printf(&server, "Started listening on port %s", SERVER_PORT);

  while (!server.has_started ||
         atomic_load_explicit(&server.num_unfinished,
                              memory_order_acquire) > 0) {
    if ((status = server_accept(&server)) != SUCCESS)
      break;
  }

  close(server.socket);

  pthread_mutex_lock(&server.threads_lock);
  server.should_close = true;
  pthread_mutex_unlock(&server.threads_lock);

  if (backup)
    pthread_join(backup_thread, NULL);

  server_release(&server);

  if (!server.has_started ||
      atomic_load_explicit(&server.num_unfinished, memory_order_acquire) > 0)
    return SYSTEM_ERROR;
  else
    return SUCCESS;
}

static error_code server_init(server *server) {
  atomic_init(&server->has_layout, false);
  server->has_started = false;

  server->owners      = NULL;
  server->subproblems = NULL;

  atomic_init(&server->num_unfinished, 0);
  atomic_init(&server->num_available, 0);

  server->num_clients = 0;

  server->should_close = false;

  server->thread_capacity = 16;
  server->num_active_threads = 0;
  server->threads = malloc(sizeof(*server->threads) * server->thread_capacity);
  if (!server->threads)
    return OUT_OF_MEMORY;

  server->client_ids = malloc(sizeof(*server->client_ids) *
                              server->thread_capacity);
  if (!server->client_ids) {
    free(server->threads);
    return OUT_OF_MEMORY;
  }

  if (pthread_mutex_init(&server->threads_lock, NULL) != 0) {
    free(server->client_ids);
    free(server->threads);
    return SYSTEM_ERROR;
  }

  server->solutions_fd = -1;

  return SUCCESS;
}

static void server_release(server *server) {
  for (uint32_t i = 0; i < server->num_active_threads; i++)
    pthread_join(server->threads[i], NULL);

  if (server->solutions_fd != -1) {
    close(server->solutions_fd);
    pthread_mutex_destroy(&server->solutions_lock);
  }

  pthread_mutex_destroy(&server->threads_lock);

  free(server->client_ids);
  free(server->threads);

  free(server->owners);
  free(server->subproblems);
}

static
error_code server_open(server *server) {
  struct addrinfo hints;
  memset(&hints, 0, sizeof hints);
  hints.ai_family   = AF_UNSPEC;
  hints.ai_socktype = SOCK_STREAM;
  hints.ai_flags    = AI_PASSIVE;

  struct addrinfo *server_info;

  int result = 0;
  if ((result = getaddrinfo(NULL, SERVER_PORT, &hints, &server_info)) != 0) {
    server_printf(server, "getaddrinfo failed: %s", gai_strerror(errno));
    return IO_ERROR;
  }

  struct addrinfo *it;

  for (it = server_info; it != NULL; it = it->ai_next) {
    server->socket = socket(it->ai_family, it->ai_socktype, it->ai_protocol);
    if (server->socket == -1) continue;

    int one = 1;
    if (setsockopt(server->socket, SOL_SOCKET, SO_REUSEADDR,
                   &one, sizeof(int)) == -1) {
      close(server->socket);
      continue;
    }

    if (bind(server->socket, it->ai_addr, it->ai_addrlen) == -1)
      close(server->socket);
    else
      break;
  }

  freeaddrinfo(server_info);

  if (it == NULL) {
    server_printf(server, "Failed to bind server");
    return IO_ERROR;
  }

  error_code status;
  if ((status = configure_nonblock(server->socket))) {
    server_printf(server, "Failed to make server non-blocking: %s",
                  strerror(errno));
    close(server->socket);
    return status;
  }

  if (listen(server->socket, 16) == -1) {
    server_printf(server, "Failed to start listening: %s",
                  strerror(errno));
    close(server->socket);
    return IO_ERROR;
  }

  return SUCCESS;
}

static error_code server_accept(server *server) {
  struct timeval timeout;
  timeout.tv_sec  = 1;
  timeout.tv_usec = 0;

  fd_set set;
  FD_ZERO(&set);
  FD_SET(server->socket, &set);

  int result = select(server->socket + 1, &set, NULL, NULL, &timeout);
  if (result == -1) {
    server_printf(server, "Failed to receive a client: %s",
                  strerror(errno));
    return SYSTEM_ERROR;
  }

  else if (result == 0)
    return SUCCESS;
  else {
    server_printf(server, "Detected a connection attempt");

    struct sockaddr_storage client_addr;
    socklen_t client_size = sizeof(client_addr);

    int client_fd = accept(server->socket, (struct sockaddr*)&client_addr,
                           &client_size);
    if (client_fd == -1) {
      server_printf(server, "Failed to accept client");
      return SUCCESS;
    }

    server_printf(server, "Accepted new client");

    client_parameters param;
    param.server = server;
    param.socket = client_fd;
    param.addr = client_addr;
    if (spawn_client_thread(server, &param) != SUCCESS)
      close(client_fd);
  }

  return SUCCESS;
}

static error_code spawn_backup_thread(server *server, pthread_t *backup_thread,
                                      server_backup_info *backup,
                                      backup_parameters *params) {
  snprintf(params->temp_path, 1024, "%s/temp.data", backup->directory);
  snprintf(params->bitmap_path, 1024, "%s/bitmap.data",
           backup->directory);
  snprintf(params->subproblems_path, 1024, "%s/subproblems.data",
           backup->directory);
  snprintf(params->solutions_path, 1024, "%s/solutions.data",
           backup->directory);

  if (pthread_mutex_init(&server->solutions_lock, NULL) != 0)  {
    server_printf(server, "Failed to create solutions mutex");
    return SYSTEM_ERROR;
  }

  int flags = O_WRONLY | O_CREAT;
  if (backup->try_restore) flags |= O_APPEND;
  server->solutions_fd = open(params->solutions_path, flags, 0644);
  if (server->solutions_fd == -1) {
    server_printf(server, "Failed to create solutions file");
    pthread_mutex_destroy(&server->solutions_lock);
    return SYSTEM_ERROR;
  }

  if (pthread_create(backup_thread, NULL, backup_handler, params) != 0) {
    server_printf(server, "Failed to create backup thread");
    return SYSTEM_ERROR;
  }

  return SUCCESS;
}

static error_code spawn_client_thread(
  server *server, client_parameters *params) {
  client_parameters *param_copy = malloc(sizeof(*params));
  if (!param_copy)
    return OUT_OF_MEMORY;

  memcpy(param_copy, params, sizeof(*params));
  param_copy->id = server->num_clients + 1;

  pthread_t thread;
  if (pthread_create(&thread, NULL, client_handler, param_copy) != 0) {
    free(param_copy);
    return SYSTEM_ERROR;
  }

  error_code status = SUCCESS;

  pthread_mutex_lock(&server->threads_lock);
  if (server->num_active_threads == server->thread_capacity) {
    uint32_t new_capacity = server->thread_capacity * 3 / 2;
    uint64_t *ids = realloc(server->client_ids, sizeof(*ids) * new_capacity);
    if (!ids)
      status = OUT_OF_MEMORY;
    else {
      server->client_ids = ids;

      pthread_t *threads = realloc(
        server->threads, sizeof(*threads) * new_capacity);
      if (!threads)
        status = OUT_OF_MEMORY;
      else {
        server->threads = threads;
        server->thread_capacity = new_capacity;
      }
    }
  }

  if (status == SUCCESS) {
    server->threads[server->num_active_threads] = thread;
    server->client_ids[server->num_active_threads] = server->num_clients + 1;
    server->num_active_threads++;
  }
  else
    pthread_detach(thread);
  pthread_mutex_unlock(&server->threads_lock);

  server->num_clients++;
  return SUCCESS;
}

static
bool server_waiting_for_layout(void *data) {
  server *server = data;
  return !server->should_close &&
    !atomic_load_explicit(&server->has_layout, memory_order_acquire);
}

static
bool server_should_stay_open(void *data) {
  server *server = data;
  return !server->should_close;
}

/**
 * Fills @p buffer with a C-string used to identify a client.
 *
 * @param [out] buffer
 * @param len Length of buffer, including any space required to store the
 *   NUL terminator.
 * @param addr
 */
static error_code get_client_name(
  server *server,
  char *buffer, size_t len,
  int socket, struct sockaddr_storage *addr);

/**
 * Waits until a subproblem layout and a set of subproblems have been received.
 *
 * If the client represented by @p socket is a master client, this function
 * reads the memory layout from it as well as its initial set of subproblems.
 *
 * @param [out] layout
 */
static error_code wait_for_layout(
  server *server, const char *client_name,
  int socket, subproblem_layout *layout);

/**
 * Reads a memory layout and a set of subproblems from client.
 *
 * This function should only be called by one client at a time. This can be
 * guaranteed using the @ref server::has_layout flag.
 *
 * @param [out] layout
 */
static error_code expect_layout(
  server *server, const char *client_name,
  int socket, subproblem_layout *layout);

/**
 * Allocates memory required to receive subproblems from a client.
 *
 * After this function returns, the buffer pointed by @p buffer should be freed
 * by the user.
 *
 * @param [out] buffer Pointer to a buffer.
 */
static error_code alloc_for_layout(
  server *server, const char *client_name,
  const subproblem_layout *layout, void **buffer,
  size_t capacity);

/**
 * Reads a set
 */
static error_code receive_initial_subproblems(
  server *server, const char *client_name,
  int socket,
  const subproblem_layout *layout, void *buffer);

/**
 * Hook called after receiving a @ref TAG_SOLUTION command from the client.
 *
 * This reads the solution from the socket and calls the solution callback.
 */
static error_code handle_solution(
  server *server, const char *client_name,
  int socket, void *buffer);

/**
 * Hook called when receiving a subproblem from the client.
 *
 * This reads the subproblem and stores it in a buffer.
 */
static error_code handle_subproblem(
  server *server, const char *client_name,
  int socket);

/**
 * Sends a subproblem layout the client.
 */
static error_code send_layout(
  server *server, const char *client_name,
  int socket);

/**
 * Initilazes a data structure used to represent a client.
 *
 * Calls to this function should be matched with calls to @ref
 * release_client_data.
 */
static error_code init_client_data(
  client_data *client,
  server *server, const char *client_name,
  uint64_t client_id, int socket);

/**
 * Releases resources created by @ref init_client_data.
 */
static void release_client_data(client_data *client);

/**
 * Releases memory used by the thread used to process a client.
 */
static void release_client_thread(server *server, uint64_t id);

static error_code process_commands(server *server, client_data *client);

/**
 * Resets the owner of all subproblems which have been sent to the client but
 * for which no finished messages have been received.
 */
static void release_unfinished_problems(server *server, client_data *client);

/**
 * Updates internal data structure to reflect the fact that @p client finished
 * to explore subproblem @p id.
 */
static error_code handle_finished(server *server, client_data *client,
                                  uint64_t id);

/**
 * Processes a request from @p client to send @p num_requested subproblems.
 *
 * This will select up to the requested number of subproblems and send them to
 * the client.
 */
static error_code handle_request(server *server, client_data *client,
                                 uint64_t num_requested);

/**
 * Makes sure the ID buffer of client can store at least @p n integers.
 */
static error_code reserve_ids(client_data *client, uint64_t n);

/**
 * Selects a list of subproblems for @p client to explore.
 *
 * Their IDs will be stored in the ID buffer of the client.
 *
 * @param num_requested Maximum number of problems to select
 * @param [out] num_selected Number of selected subproblems.
 */
static void select_subproblems(server *server, client_data *client,
                               uint64_t num_requested,
                               uint64_t *num_selected);

/**
 * Sends @p num_selected problems to @p client. The IDs of those problems are
 * stored in @ref client::id_buffer.
 */
static error_code send_subproblems(
  server *server, client_data *client, uint64_t num_selected);

static void *client_handler(void *thread_arg) {
  client_parameters *params = thread_arg;

  char client_name[512];
  if (get_client_name(params->server, client_name, sizeof(client_name),
                      params->socket, &params->addr) != SUCCESS)
    goto fail;

  server_printf(params->server, "Processing client %s", client_name);

  if (configure_nonblock(params->socket) != SUCCESS) {
    server_printf(params->server, "[%s] Failed to make client socket "
                  "non-blocking\n", client_name);
    goto fail;
  }

  configure_keepalive(params->socket);

  subproblem_layout layout;
  if (wait_for_layout(params->server, client_name,
                      params->socket, &layout) != SUCCESS) {
    goto fail;
  }

  if (send_layout(params->server, client_name, params->socket) != SUCCESS)
    goto fail;

  client_data client;
  if (init_client_data(&client, params->server, client_name,
                       params->id, params->socket) != SUCCESS) {
    goto fail;
  }

  if (process_commands(params->server, &client) != SUCCESS) {
    goto release;
  }

release:
  release_unfinished_problems(params->server, &client);
  release_client_data(&client);

fail:
  server_printf(params->server, "[%s] Cleaning up connection", client_name);

  close(params->socket);
  release_client_thread(params->server, params->id);
  free(params);
  return NULL;
}

static error_code get_client_name(
  server *server,
  char *buffer, size_t len,
  int socket, struct sockaddr_storage *addr) {
  void *in_addr;
  if (addr->ss_family == AF_INET)
    in_addr = &(((struct sockaddr_in*)addr)->sin_addr);
  else
    in_addr = &(((struct sockaddr_in6*)addr)->sin6_addr);

  char client_login[256];
  if (!read_name(server_should_stay_open, server,
                 socket, client_login, sizeof(client_login))) {
    server_printf(server, "Failed to read user name\n");
    return PROTOCOL_ERROR;
  }

  char client_addr[256];
  inet_ntop(addr->ss_family, in_addr, client_addr, sizeof(client_addr));

  snprintf(buffer, len, "%s@%s", client_login, client_addr);
  return SUCCESS;
}

static error_code wait_for_layout(
  server *server, const char *client_name,
  int socket, subproblem_layout *layout) {
  message_tag tag;
  uint64_t arg;
  bool received_command = false;
  if (!(received_command =
        read_command(server_waiting_for_layout, server,
                     socket, &tag, &arg)) &&
      server_waiting_for_layout(server)) {
    return SYSTEM_ERROR;
  }

  if (received_command) {
    if (tag != TAG_LAYOUT)
      server_printf(server, "[%s] Invalid first command: %d", client_name, tag);
    else {
      bool expected = false;
      if (!atomic_compare_exchange_strong_explicit(
            &server->has_layout, &expected, true,
            memory_order_acq_rel,
            memory_order_acquire)) {
        server_printf(server, "[%s] Lost race to become master client",
                      client_name);
        return SYSTEM_ERROR;
      }

      if (expect_layout(server, client_name, socket, layout) != SUCCESS) {
        atomic_store_explicit(&server->has_layout, false, memory_order_release);
        return SYSTEM_ERROR;
      }
    }
  }
  else
    server_printf(server, "[%s] Waiting for memory layout", client_name);

  while (server_should_stay_open(server) && !server->has_started)
    usleep(1000 * 100);

  *layout = server->layout;
  return SUCCESS;
}

static error_code expect_layout(
  server *server, const char *client_name,
  int socket, subproblem_layout *layout) {
  server_printf(server, "[%s] Listening for memory layout", client_name);

  uint32_t num_vertices, num_quads, num_hexes,
    num_symmetries, num_boundary_vertices, max_no_goods_size;
  uint8_t flags;
  if (!read_layout(server_should_stay_open, server,
                   socket, &num_vertices, &num_quads, &num_hexes,
                   &num_symmetries, &num_boundary_vertices,
                   &max_no_goods_size, &flags)) {
    server_printf(server, "[%s] Failed to get layout", client_name);
    return SYSTEM_ERROR;
  }

  server_printf(server, "[%s] Got layout (V=%u, Q=%u, H=%u)",
                client_name, num_vertices, num_quads, num_hexes);

  compute_subproblem_layout(layout, num_vertices, num_quads, num_hexes,
                            num_symmetries, num_boundary_vertices,
                            max_no_goods_size, flags);
  server->layout = *layout;

  void *buffer;
  if (alloc_for_layout(server, client_name, layout, &buffer,
                       SERVER_CAPACITY) != SUCCESS)
    return OUT_OF_MEMORY;

  server_printf(server, "[%s] Waiting for initial subproblems", client_name);

  if (receive_initial_subproblems(server, client_name, socket,
                                  layout, buffer) != SUCCESS) {
    free(buffer);
    free(server->subproblems);
    free(server->owners);
    return SYSTEM_ERROR;
  }

  free(buffer);
  return SUCCESS;
}


static error_code alloc_for_layout(
  server *server, const char *client_name,
  const subproblem_layout *layout, void **buffer,
  size_t capacity) {
  void *data_buffer = malloc(layout->packet_size);
  if (!data_buffer) {
    server_printf(server, "[%s] Failed to allocate data buffer "
                  "of size %zu", client_name, layout->packet_size);
    return OUT_OF_MEMORY;
  }
  else {
    server_printf(server, "[%s] Allocated data buffer of size %zu",
                  client_name, layout->packet_size);
  }

  *buffer = data_buffer;

  server->subproblems = malloc(capacity * layout->packet_size);
  if (!server->subproblems) {
    server_printf(server, "[%s] Failed to allocate solver pool "
                  "of size %u", client_name, capacity);
    free(data_buffer);
    return OUT_OF_MEMORY;
  }

  server->owners = malloc(sizeof(*server->owners) * capacity);
  if (!server->owners) {
    free(server->subproblems);
    free(data_buffer);
    return OUT_OF_MEMORY;
  }

  server->num_subproblems = 0;
  server->capacity = capacity;

  atomic_store_explicit(&server->num_unfinished, 0, memory_order_release);

  return SUCCESS;
}

static error_code receive_initial_subproblems(
  server *server, const char *client_name,
  int socket,
  const subproblem_layout *layout, void *buffer) {
  bool received_command;

  message_tag tag;
  uint64_t arg;

  while ((received_command =
          read_command(server_should_stay_open, server,
                       socket, &tag, &arg)) &&
         tag != TAG_EMPTY) {
    switch (tag) {
    case TAG_SOLUTION:
      if (handle_solution(server, client_name, socket, buffer) != SUCCESS)
        return SYSTEM_ERROR;
      break;
    case TAG_SUBPROBLEM:
      if (handle_subproblem(server, client_name, socket) != SUCCESS)
        return SYSTEM_ERROR;
      break;
    default:
      server_printf(server, "[%s] Invalid command during initial problem "
                    "reception: %d", client_name, tag);
      return SYSTEM_ERROR;
    }
  }

  if (!received_command) {
    server_printf(server, "[%s] Initial subproblem reception failed",
                  client_name);
    return SYSTEM_ERROR;
  }
  else {
    server_printf(server, "[%s] Got %u subproblems",
                  client_name, server->num_subproblems);

    atomic_store_explicit(&server->num_unfinished, server->num_subproblems,
                          memory_order_release);
    atomic_store_explicit(&server->num_available, server->num_subproblems,
                          memory_order_release);

    server->has_started = true;

    return SUCCESS;
  }
}

static error_code handle_solution(
  server *server, const char *client_name,
  int socket, void *buffer) {
  vertex_index *vertices;
  hex_index num_hexes;
  if (!read_solution(
        server_should_stay_open, server,
        socket, buffer, &vertices, &num_hexes,
        server->layout.num_hexes)) {
    server_printf(server, "[%s] Failed to read solution", client_name);
    return SYSTEM_ERROR;
  }

  server_printf(server, "[%s] Received a solution!", client_name);
  server->cb(client_name, vertices, num_hexes, server->user_data);

  if (server->solutions_fd != -1) {
    pthread_mutex_lock(&server->solutions_lock);
    if (!write_solution(server_should_stay_open, server, server->solutions_fd,
                        buffer, vertices, num_hexes)) {
      server_printf(server, "[BACKUP] Failed to store a solution!");
      close(server->solutions_fd);
      server->solutions_fd = -1;
    }

    pthread_mutex_unlock(&server->solutions_lock);
    if (server->solutions_fd == -1)
      pthread_mutex_destroy(&server->solutions_lock);
  }

  return SUCCESS;
}

static error_code handle_subproblem(
  server *server, const char *client_name,
  int socket) {
  if (server->num_subproblems == server->capacity) {
    uint64_t new_capa = (server->capacity * 3) / 2;
    uint8_t *buffer = realloc(server->subproblems, new_capa *
                              server->layout.packet_size);
    if (!buffer)
      return OUT_OF_MEMORY;
    server->subproblems = buffer;

    atomic_uint_fast32_t *owners =
      realloc(server->owners, sizeof(*owners) * new_capa);
    if (!owners)
      return OUT_OF_MEMORY;

    server->owners = owners;

    server->capacity = new_capa;
  }

  uint8_t *read_buffer = server->subproblems +
    server->layout.packet_size * (size_t)server->num_subproblems;

  if (!read_raw_subproblem(server_should_stay_open, server,
                           socket, read_buffer, &server->layout)) {
    return SYSTEM_ERROR;
  }

  atomic_init(&server->owners[server->num_subproblems++], 0);

  atomic_fetch_add_explicit(&server->num_unfinished, 1, memory_order_release);
  return SUCCESS;
}

static error_code send_layout(
  server *server, const char *client_name,
  int socket) {
  server_printf(server, "[%s] Sending memory layout",
                client_name);

  if (!write_command(server_should_stay_open, server,
                     socket, TAG_LAYOUT, 0)) {
    server_printf(server, "[%s] Failed to send layout", client_name);
    return SYSTEM_ERROR;
  }

  if (!write_layout(server_should_stay_open, server,
                    socket, server->layout.num_vertices,
                    server->layout.num_quads, server->layout.num_hexes,
                    server->layout.num_symmetries,
                    server->layout.num_boundary_vertices,
                    server->layout.max_no_goods_size,
                    server->layout.flags)) {
    server_printf(server, "[%s] Failed to send layout", client_name);
    return SYSTEM_ERROR;
  }

  return SUCCESS;
}

static error_code init_client_data(
  client_data *client,
  server *server, const char *client_name,
  uint64_t client_id, int socket) {
  client->id     = client_id;
  client->name   = client_name;
  client->socket = socket;

  client->cursor_id = 0;

  client->buffer = malloc(server->layout.packet_size);
  if (!client->buffer) return OUT_OF_MEMORY;

  if (sparse_set_init(&client->exploring, server->num_subproblems) != SUCCESS) {
    free(client->buffer);
    return OUT_OF_MEMORY;
  }

  client->id_buffer   = NULL;
  client->id_capacity = 0;

  client->asleep = false;

  return SUCCESS;
}

static void release_client_data(client_data *client) {
  free(client->id_buffer);
  sparse_set_release(&client->exploring);
  free(client->buffer);
}

static void release_client_thread(server *server, uint64_t id) {
  pthread_mutex_lock(&server->threads_lock);
  if (!server->should_close) {
    uint32_t position;
    for (position = 0; position < server->num_active_threads &&
           server->client_ids[position] != id; position++);

    if (position < server->num_active_threads) {
      pthread_detach(server->threads[position]);
      server->threads[position] =
        server->threads[server->num_active_threads - 1];
      server->client_ids[position] =
        server->client_ids[server->num_active_threads - 1];

      server->num_active_threads--;
    }
  }
  pthread_mutex_unlock(&server->threads_lock);
}

typedef struct connection_data {
  server *server;
  client_data *client;
} connection_data;

static bool should_wake_up_client(server *server, client_data *client) {
  return client->asleep &&
    atomic_load_explicit(&server->num_available, memory_order_acquire) != 0;
}

static bool should_wait_or_wake_up(void *data) {
  connection_data *connection = (connection_data*)data;
  return server_should_stay_open(connection->server) &&
    !should_wake_up_client(connection->server, connection->client);
}

static error_code process_commands(server *server, client_data *client) {
  message_tag tag;
  uint64_t arg;

  error_code status = SUCCESS;

  while (true) {
    connection_data connection = {server, client};
    bool received = read_command(should_wait_or_wake_up, &connection,
                                 client->socket, &tag, &arg);
    if (received) {
      switch (tag) {
      case TAG_CLOSE:
        server_printf(server, "[%s] Manually closing connection", client->name);
        goto finish;
      case TAG_FINISHED:
        if ((status = handle_finished(server, client, arg)) != SUCCESS)
          goto finish;
        break;
      case TAG_SOLUTION:
        if ((status = handle_solution(server, client->name, client->socket,
                                      client->buffer)) != SUCCESS)
          goto finish;
        break;
      case TAG_REQUEST:
        if ((status = handle_request(server, client, arg)) != SUCCESS)
          goto finish;
        break;
      default:
        server_printf(server, "[%s] Invalid command: %d", client->name, tag);
        status = SYSTEM_ERROR;
        goto finish;
      }
    }
    else if (should_wake_up_client(server, client)) {
      if (!write_command(server_should_stay_open, server,
                         client->socket, TAG_WAKEUP, 0)) {
        status = SYSTEM_ERROR;
        break;
      }

      client->asleep = false;
    }
    else
      break;
  }

finish:
  return status;
}

static void release_unfinished_problems(
  server *server, client_data *client) {
  if (client->exploring.size != 0) {
    server_printf(server,
                  "[%s] Re-inserting %u subproblems into the queue",
                  client->name, client->exploring.size);
  }

  for (uint32_t i = 0; i < client->exploring.size; i++) {
    uint32_t id = client->exploring.data[i];
    atomic_store_explicit(&server->owners[id], 0, memory_order_release);
    atomic_fetch_add_explicit(&server->num_available, 1, memory_order_acq_rel);
  }
}

static error_code handle_finished(server *server, client_data *client,
                                  uint64_t id) {
  if (id >= server->num_subproblems) {
    server_printf(server, "[%s] Requested to release unowned problem %zu",
                  client->name, id);
    return SYSTEM_ERROR;
  }

  uint64_t owner = atomic_load_explicit(&server->owners[id],
                                        memory_order_acquire);

  if (owner != client->id) {
    server_printf(server, "[%s] Requested to release unowned problem %zu",
                  client->name, id);
    return SYSTEM_ERROR;
  }

  atomic_store_explicit(&server->owners[id], OWNER_DONE,
                        memory_order_release);
  atomic_store_explicit(&server->bitmap_dirty, true,
                        memory_order_release);

  sparse_set_remove(&client->exploring, id);
  uint32_t still_unfinished = atomic_fetch_sub_explicit(
    &server->num_unfinished, 1, memory_order_acq_rel) - 1;
  if (still_unfinished % 100 == 0) {
    server->progress_cb(server->num_subproblems - still_unfinished,
                        server->num_subproblems, server->progress_data);
  }

  return SUCCESS;
}

static error_code handle_request(
  server *server, client_data *client, uint64_t num_requested) {
  server_printf(server,
                "[%s] Requested %zu subproblems (num left: %u, "
                "num exploring: %u)",
                client->name, num_requested,
                atomic_load_explicit(&server->num_unfinished,
                                     memory_order_acquire),
                client->exploring.size);
  reserve_ids(client, num_requested);

  if (client->id_capacity < num_requested) {
    num_requested = client->id_capacity;
    server_printf(server, "[%s] Client requested too many subproblems, "
                  "only %zu will be selected", client->name, num_requested);
  }

  uint64_t num_selected;
  select_subproblems(server, client, num_requested, &num_selected);

  error_code status;
  if ((status = send_subproblems(server, client, num_selected)) != SUCCESS) {
    server_printf(server, "[%s] Failed to send %zu subproblems",
                  client->name, num_selected);
  }
  else {
    server_printf(server, "[%s] Sent %zu subproblems",
                  client->name, num_selected);

    if (num_selected == 0) {
      server_printf(server, "[%s] No more available subproblems, going to sleep",
                    client->name);
      client->asleep = true;
    }
  }

  return status;
}

static error_code reserve_ids(client_data *client, uint64_t n) {
  if (client->id_capacity < n) {
    uint32_t *buffer = realloc(client->id_buffer, n * sizeof(*buffer));
    if (buffer) {
      client->id_buffer   = buffer;
      client->id_capacity = n;
    }
    else
      return OUT_OF_MEMORY;
  }

  return SUCCESS;
}

static void select_subproblems(
  server *server, client_data *client,
  uint64_t num_requested, uint64_t *num_selected) {
  *num_selected = 0;

  if (server->num_subproblems != 0) {
    uint32_t old_id = client->cursor_id;
    do {
      uint64_t expected = 0;
      if (atomic_compare_exchange_strong_explicit(
            &server->owners[client->cursor_id],
            &expected, client->id,
            memory_order_acq_rel,
            memory_order_acquire)) {
        client->id_buffer[(*num_selected)++] = client->cursor_id;
        sparse_set_insert(&client->exploring, client->cursor_id);

        atomic_fetch_sub_explicit(&server->num_available, 1,
                                  memory_order_acq_rel);
      }

      client->cursor_id++;
      if (client->cursor_id == server->num_subproblems)
        client->cursor_id = 0;
    } while (old_id != client->cursor_id && *num_selected < num_requested);
  }
}

static error_code send_subproblems(server *server, client_data *client,
                                   uint64_t num_selected) {
  for (uint64_t i = 0; i < num_selected; i++) {
    uint64_t id = client->id_buffer[i];

    if (!write_command(server_should_stay_open, server,
                       client->socket, TAG_SUBPROBLEM, id)) {
      return SYSTEM_ERROR;
    }

    uint8_t *buffer = server->subproblems + server->layout.packet_size * id;

    if (!write_raw_subproblem(server_should_stay_open, server,
                              client->socket, buffer,
                              &server->layout))
      return SYSTEM_ERROR;
  }

  if (!write_command(server_should_stay_open, server,
                     client->socket, TAG_EMPTY, 0))
    return SYSTEM_ERROR;

  return SUCCESS;
}

#define BACKUP_SLEEP_TIME 5

static bool always(void *arg) { return true; }

static void backup_restore(backup_parameters *backup);

/**
 * Stores all subproblems collected by the server to a file.
 */
static error_code backup_write_subproblems(backup_parameters *backup);

/**
 * Saves progress made by the algorithm to a file.
 */
static error_code backup_write_bitmap(backup_parameters *backup,
                                      uint64_t *buffer, size_t buffer_size);

/**
 * As long as the server is running, regularly updates the backup.
 */
static void backup_update_loop(backup_parameters *backup);

static void *backup_handler(void *arg) {
  backup_parameters *backup = arg;

  if (backup->try_restore) {
    backup_restore(backup);
    backup_update_loop(backup);
    return NULL;
  }

  while (server_should_stay_open(backup->server) &&
         !backup->server->has_started)
    sleep(BACKUP_SLEEP_TIME);

  if (!server_should_stay_open(backup->server))
    return NULL;

  if (backup_write_subproblems(backup) != SUCCESS) {
    server_printf(backup->server, "[BACKUP] Failed to write subproblems, "
                  "won't be able to restore progress!");
    return NULL;
  }

  backup_update_loop(backup);

  return NULL;
}

static error_code backup_write_subproblems(backup_parameters *backup) {
  int fd = open(backup->temp_path, O_WRONLY | O_CREAT, 0644);
  if (fd == -1)
    return IO_ERROR;

  uint8_t buffer[sizeof(uint32_t)];
  encode_uint32(buffer, backup->server->num_subproblems);

  if (write(fd, buffer, sizeof(buffer)) != (ssize_t)sizeof(buffer)) {
    close(fd);
    return IO_ERROR;
  }

  if (!write_layout(always, NULL, fd,
                    backup->server->layout.num_vertices,
                    backup->server->layout.num_quads,
                    backup->server->layout.num_hexes,
                    backup->server->layout.num_symmetries,
                    backup->server->layout.num_boundary_vertices,
                    backup->server->layout.max_no_goods_size,
                    backup->server->layout.flags)) {
    close(fd);
    return IO_ERROR;
  }

  size_t buf_size = backup->server->layout.packet_size *
    (size_t)backup->server->num_subproblems;
  ssize_t n = write(fd, backup->server->subproblems, buf_size);
  if (n < 0 || (size_t)n != buf_size) {
    close(fd);
    return IO_ERROR;
  }

  close(fd);

  if (rename(backup->temp_path, backup->subproblems_path) != 0)
    return IO_ERROR;

  return SUCCESS;
}

static void backup_update_loop(backup_parameters *backup) {
  size_t buffer_size = sizeof(uint64_t) *
    ((size_t)backup->server->num_subproblems + 63) / 64;
  void *buffer = malloc(buffer_size);

  if (!buffer) {
    server_printf(backup->server, "[BACKUP] Failed to allocate bitmap, "
                  "won't be able to restore progress!");
    return;
  }

  if (backup_write_bitmap(backup, buffer, buffer_size) != SUCCESS) {
    server_printf(backup->server, "[BACKUP] Failed to write bitmap, "
                  "may not be able to restore progress!");
  }

  while (server_should_stay_open(backup->server)) {
    bool expected = true;
    if (atomic_compare_exchange_strong_explicit(
          &backup->server->bitmap_dirty,
          &expected, false,
          memory_order_acq_rel,
          memory_order_acquire)) {
      if (backup_write_bitmap(backup, buffer, buffer_size) != SUCCESS) {
        server_printf(backup->server, "[BACKUP] Failed to write bitmap, "
                      "may not be able to restore progress!");
        atomic_store_explicit(&backup->server->bitmap_dirty, true,
                              memory_order_release);
      }
    }

    sleep(BACKUP_SLEEP_TIME);
  }

  backup_write_bitmap(backup, buffer, buffer_size);

  free(buffer);
}

static error_code backup_write_bitmap(backup_parameters *backup,
                                      uint64_t *buffer, size_t buffer_size) {
  memset(buffer, 0, buffer_size);

  for (uint32_t i = 0; i < backup->server->num_subproblems; i++) {
    if (atomic_load_explicit(&backup->server->owners[i],
                             memory_order_acquire) == OWNER_DONE) {
      buffer[i / 64] |= (1ul << (i % 64));
    }
  }

  for (size_t i = 0; i < buffer_size / sizeof(uint64_t); i++)
    encode_uint64(&buffer[i], buffer[i]);

  int fd = open(backup->temp_path, O_WRONLY | O_CREAT, 0644);
  if (fd == -1) return IO_ERROR;

  ssize_t n = write(fd, buffer, buffer_size);
  if (n < 0 || (size_t)n != buffer_size) {
    close(fd);
    return IO_ERROR;
  }

  close(fd);

  if (rename(backup->temp_path, backup->bitmap_path) != 0)
    return IO_ERROR;

  server_printf(backup->server, "[BACKUP] Progress stored to disk");

  return SUCCESS;
}

static error_code backup_restore_subproblems(backup_parameters *backup, void **buffer);
static error_code backup_restore_progress(backup_parameters *backup);
static error_code backup_restore_solutions(backup_parameters *backup, void *buffer);

static void backup_restore(backup_parameters *backup) {
  bool expected = false;
  if (!atomic_compare_exchange_strong_explicit(
        &backup->server->has_layout, &expected, true,
        memory_order_acq_rel,
        memory_order_acquire)) {
    server_printf(backup->server, "[BACKUP] Server is already listening for "
                  "a subproblem layout");
    return;
  }

  void *buffer;
  if (backup_restore_subproblems(backup, &buffer) != SUCCESS) {
    atomic_store_explicit(&backup->server->has_layout, false,
                          memory_order_release);
    return;
  }

  if (backup_restore_progress(backup) != SUCCESS) {
    free(buffer);
    backup->server->has_started = true;
    return;
  }

  if (backup_restore_solutions(backup, buffer) != SUCCESS) {
    free(buffer);
  }

  backup->server->has_started = true;
}

static error_code backup_restore_subproblems(backup_parameters *backup,
                                             void **buffer) {
  int fd = open(backup->subproblems_path, O_RDONLY);
  if (fd == -1) {
    server_printf(backup->server, "[BACKUP] Failed to open old "
                  "subproblems file");
    return IO_ERROR;
  }

  uint32_t num_subproblems;

  {
    uint8_t buffer[sizeof(uint32_t)];
    if (read(fd, buffer, sizeof(buffer)) != (ssize_t)sizeof(buffer)) {
      server_printf(backup->server, "[BACKUP] Failed to read old subproblems!");
      close(fd);
      return IO_ERROR;
    }

    num_subproblems = decode_uint32(buffer);
  }

  uint32_t num_vertices, num_quads, num_hexes, num_symmetries,
    num_boundary_vertices, max_no_goods_size;
  uint8_t flags;
  if (!read_layout(always, NULL, fd, &num_vertices, &num_quads, &num_hexes,
                   &num_symmetries, &num_boundary_vertices,
                   &max_no_goods_size, &flags)) {
    server_printf(backup->server, "[BACKUP] Failed to get layout");
    close(fd);
    return IO_ERROR;
  }

  compute_subproblem_layout(&backup->server->layout,
                            num_vertices, num_quads, num_hexes,
                            num_symmetries, num_boundary_vertices,
                            max_no_goods_size, flags);

  if (alloc_for_layout(backup->server, "BACKUP",
                       &backup->server->layout, buffer,
                       num_subproblems) != SUCCESS) {
    close(fd);
    return IO_ERROR;
  }

  size_t n_read = backup->server->layout.packet_size *
    (size_t)num_subproblems;

  if (read(fd, backup->server->subproblems, n_read) != (ssize_t)n_read) {
    server_printf(backup->server, "[BACKUP] Failed to read!");

    free(*buffer);
    free(backup->server->subproblems);
    free(backup->server->owners);

    close(fd);

    return IO_ERROR;
  }

  close(fd);

  backup->server->num_subproblems = num_subproblems;

  atomic_store_explicit(&backup->server->num_unfinished, num_subproblems,
                        memory_order_release);
  atomic_store_explicit(&backup->server->num_available, num_subproblems,
                        memory_order_release);

  for (uint32_t i = 0; i < num_subproblems; i++)
    atomic_init(&backup->server->owners[i], 0);

  return SUCCESS;
}

static error_code backup_restore_progress(backup_parameters *backup) {
  int fd = open(backup->bitmap_path, O_RDONLY);
  if (fd == -1) {
    server_printf(backup->server, "[BACKUP] Failed to restore solved subproblems!");
    close(fd);
    return IO_ERROR;
  }

  size_t bitmap_size = sizeof(uint64_t) *
    (((size_t)backup->server->num_subproblems + 63) / 64);
  uint64_t *bitmap = malloc(bitmap_size);

  if (!bitmap) {
    server_printf(backup->server, "[BACKUP] Failed to allocate bitmap, "
                  "won't be able to restore progress!");
    close(fd);
    return OUT_OF_MEMORY;
  }

  if (read(fd, bitmap, bitmap_size) != (ssize_t)bitmap_size) {
    server_printf(backup->server, "[BACKUP] Failed to read bitmap, "
                  "won't be able to restore progress!");
    close(fd);
    return OUT_OF_MEMORY;
  }

  for (size_t i = 0; i < bitmap_size / sizeof(uint64_t); i++)
    bitmap[i] = decode_uint64(&bitmap[i]);

  uint32_t num_finished = 0;
  for (uint32_t i = 0; i < backup->server->num_subproblems; i++) {
    if (bitmap[i / 64] & (1ul << (i % 64))) {
      atomic_store_explicit(&backup->server->owners[i], OWNER_DONE,
                            memory_order_release);
      num_finished++;
    }
  }

  atomic_fetch_sub_explicit(&backup->server->num_unfinished, num_finished,
                            memory_order_acq_rel);
  atomic_fetch_sub_explicit(&backup->server->num_available, num_finished,
                            memory_order_acq_rel);

  free(bitmap);
  close(fd);

  return SUCCESS;
}

static error_code backup_restore_solutions(backup_parameters *backup,
                                           void *buffer) {
  int fd = open(backup->solutions_path, O_RDONLY);
  if (fd == -1) {
    server_printf(backup->server, "[BACKUP] Failed to restore old solutions!");
    return IO_ERROR;
  }

  vertex_index *vertices;
  hex_index solution_size;
  while (read_solution(always, NULL, fd, buffer, &vertices, &solution_size,
                       backup->server->layout.num_hexes)) {
    backup->server->cb("BACKUP", vertices, solution_size,
                       backup->server->user_data);
  }

  if (read(fd, buffer, 1) != 0) {
    server_printf(backup->server, "[BACKUP] Failed to restore some of the "
                  "old solutions!");
  }

  close(fd);
  free(buffer);

  return SUCCESS;
}
